package com.yoshino.leetcode.improve40.twentysecond;

class MapSum {

    // 前缀树
    class PrefixNode {
        Integer val = 0;
        PrefixNode[] next;
        public PrefixNode() {
            next = new PrefixNode[26];
        }
    }

    PrefixNode root;
    /** Initialize your data structure here. */
    public MapSum() {
        root = new PrefixNode();
    }

    public void insert(String key, int val) {
        PrefixNode cur = root;
        for (char c : key.toCharArray()) {
            int ind = c - 'a';
            if (cur.next[ind] == null) {
                cur.next[ind] = new PrefixNode();
            }
            cur = cur.next[ind];
        }
        cur.val = val;
    }

    // 找到对应的前缀结点
    // 深度遍历后续所有可能值之和
    public int sum(String prefix) {
        PrefixNode cur = root;
        for (char c : prefix.toCharArray()) {
            if (cur.next[c - 'a'] == null) {
                return 0;
            }
            cur = cur.next[c - 'a'];
        }
        // 加上当前val，并遍历后续next
        return dfs(cur);
    }

    private int dfs(PrefixNode cur) {
        if (cur == null) {
            return 0;
        }
        int res = cur.val;
        for (int i = 0; i < 26; i++) {
            res += dfs(cur.next[i]);
        }
        return res;
    }
}

/**
 * Your MapSum object will be instantiated and called as such:
 * MapSum obj = new MapSum();
 * obj.insert(key,val);
 * int param_2 = obj.sum(prefix);
 */